from gxl_ai_utils.utils import utils_file

from wenet.utils.gxl_utils import Whisper_Utils
from wenet.utils.init_model import init_model
from wenet.utils.init_tokenizer import init_tokenizer
from wenet.utils.train_utils import check_modify_and_save_config


def do_test_model():
    configs = utils_file.load_dict_from_yaml('./conf/train_whisper_small_streaming_2.yaml')
    args = utils_file.do_dict2simpleNamespaceObj(utils_file.load_dict_from_yaml('./argparse_run.yaml'))
    utils_file.makedir_sil(args.model_dir)
    # init tokenizer
    tokenizer = init_tokenizer(configs, './data/units_en_cn.txt', './data/en_cn_bpe.model',
                               None)
    configs['vocab_size'] = tokenizer.vocab_size()
    # Do some sanity checks and save config to arsg.model_dir
    configs = check_modify_and_save_config(args, configs, tokenizer.symbol_table)
    model, _ = init_model(args, configs)
    utils_file.print_model_size(model)
    utils_file.print_model_size(model.encoder)
    utils_file.print_model_size(model.decoder)


def do_test_whisper():
    """"""
    model, _ = Whisper_Utils.load_whisper('small')
    print(model)
    num_params = sum(p.numel() for p in model.parameters())
    print('the number of model params: {:,f}M'.format(num_params / 1024 / 1024))

    encoder = model.encoder
    num_params = sum(p.numel() for p in encoder.parameters())
    print('the number of encoder params: {:,f}M'.format(num_params / 1024 / 1024))

    decoder = model.decoder
    num_params = sum(p.numel() for p in decoder.parameters())
    print('the number of decoder params: {:,f}M'.format(num_params / 1024 / 1024))
    ctc = model.ctc
    num_params = sum(p.numel() for p in ctc.parameters())
    print('the number of ctc params: {:,f}M'.format(num_params / 1024 / 1024))


if __name__ == '__main__':
    do_test_model()
